# game.py
import random

from numba import jit, int32, boolean

from config import *
from draw import draw_snake_and_food, fpsClock
from utils import load_high_score, save_high_score

# 初始化游戏板、蛇的位置和大小
board = np.zeros(FIELD_SIZE, dtype=np.int32)
snake = np.zeros(FIELD_SIZE + 1, dtype=np.int32)
snake[HEAD] = 1 * WIDTH + 1  # 初始化蛇头在(1,1)的位置
snake_size = 1  # 初始蛇的长度

# 临时变量，用于模拟蛇的移动
tmpboard = np.zeros(FIELD_SIZE, dtype=np.int32)
tmpsnake = np.zeros(FIELD_SIZE + 1, dtype=np.int32)
tmpsnake[HEAD] = 1 * WIDTH + 1
tmpsnake_size = 1

food = 4 * WIDTH + 7  # 初始食物位置
best_move = ERR  # 运动方向

score = 1  # 初始分数
high_score = 0
step_counter = 0  # 步数计数器

@jit(boolean(int32, int32, int32[:]), nopython=True)
def is_cell_free(idx, psize, psnake_array):
    """检查一个cell有没有被蛇身覆盖，没有覆盖则为free，返回true"""
    for i in range(psize):
        if psnake_array[i] == idx:
            return False
    return True

@jit(boolean(int32, int32), nopython=True)
def is_move_possible(idx, move):
    """检查某个位置idx是否可向move方向运动"""
    if move == LEFT:
        return idx % WIDTH > 1
    elif move == RIGHT:
        return idx % WIDTH < (WIDTH - 2)
    elif move == UP:
        return idx > (2 * WIDTH - 1)
    elif move == DOWN:
        return idx < (FIELD_SIZE - 2 * WIDTH)
    return False

@jit(nopython=True)
def board_reset(psnake, psize, pboard, psnake_array, food):
    """重置board"""
    for i in range(FIELD_SIZE):
        if i == food:
            pboard[i] = FOOD
        elif is_cell_free(i, psize, psnake_array):
            pboard[i] = UNDEFINED
        else:
            pboard[i] = SNAKE

@jit(boolean(int32, int32[:], int32[:], int32[:]), nopython=True)
def board_BFS(pfood, psnake, pboard, moves):
    """广度优先搜索遍历整个board，计算出board中每个非SNAKE元素到达食物的路径长度"""
    queue = [pfood]
    inqueue = [False] * FIELD_SIZE
    found = False
    while queue:
        idx = queue.pop(0)
        if inqueue[idx]:
            continue
        inqueue[idx] = True
        for move in moves:
            if is_move_possible(idx, move):
                if idx + move == psnake[HEAD]:
                    found = True
                if pboard[idx + move] < SNAKE:
                    if pboard[idx + move] > pboard[idx] + 1:
                        pboard[idx + move] = pboard[idx] + 1
                    if not inqueue[idx + move]:
                        queue.append(idx + move)
    return found

@jit(int32(int32[:], int32[:], int32[:]), nopython=True)
def choose_shortest_safe_move(psnake, pboard, moves):
    """从蛇头开始，根据board中元素值，从蛇头周围4个领域点中选择最短路径"""
    min_distance = SNAKE
    best_move = ERR
    for move in moves:
        if is_move_possible(psnake[HEAD], move) and pboard[psnake[HEAD] + move] < min_distance:
            min_distance = pboard[psnake[HEAD] + move]
            best_move = move
    return best_move

@jit(int32(int32[:], int32[:], int32[:]), nopython=True)
def choose_longest_safe_move(psnake, pboard, moves):
    """从蛇头开始，根据board中元素值，从蛇头周围4个领域点中选择最远路径"""
    max_distance = -1
    best_move = ERR
    for move in moves:
        if is_move_possible(psnake[HEAD], move) and pboard[psnake[HEAD] + move] < UNDEFINED and pboard[psnake[HEAD] + move] > max_distance:
            max_distance = pboard[psnake[HEAD] + move]
            best_move = move
    return best_move

def is_tail_inside():
    """检查是否可以追着蛇尾运动，即蛇头和蛇尾间是有路径的"""
    tmpboard[tmpsnake[tmpsnake_size - 1]] = 0
    tmpboard[food] = SNAKE
    result = board_BFS(tmpsnake[tmpsnake_size - 1], tmpsnake, tmpboard, MOVES)
    for move in MOVES:
        if is_move_possible(tmpsnake[HEAD], move) and tmpsnake[HEAD] + move == tmpsnake[tmpsnake_size - 1] and tmpsnake_size > 3:
            result = False
    return result

def follow_tail():
    """让蛇头朝着蛇尾运行一步"""
    global tmpsnake_size
    tmpsnake_size = snake_size
    tmpsnake[:] = snake[:]
    board_reset(tmpsnake, tmpsnake_size, tmpboard, tmpsnake[:tmpsnake_size], food)
    tmpboard[tmpsnake[tmpsnake_size - 1]] = FOOD
    tmpboard[food] = SNAKE
    board_BFS(tmpsnake[tmpsnake_size - 1], tmpsnake, tmpboard, MOVES)
    tmpboard[tmpsnake[tmpsnake_size - 1]] = SNAKE
    return choose_longest_safe_move(tmpsnake, tmpboard, MOVES)

def any_possible_move():
    """在各种方案都不行时，随便找一个可行的方向来走(1步)"""
    board_reset(snake, snake_size, board, snake[:snake_size], food)
    board_BFS(food, snake, board, MOVES)
    min_distance = SNAKE
    best_move = ERR
    for move in MOVES:
        if is_move_possible(snake[HEAD], move) and board[snake[HEAD] + move] < min_distance:
            min_distance = board[snake[HEAD] + move]
            best_move = move
    return best_move

@jit(nopython=True)
def shift_array(arr, size):
    """转换数组函数"""
    for i in range(size, 0, -1):
        arr[i] = arr[i - 1]

def new_food():
    """随机生成新的食物"""
    global food, snake_size
    cell_free = False
    snake_array = np.array(snake[:snake_size], dtype=np.int32)  # 使用 numpy 数组
    while not cell_free:
        w = random.randint(1, WIDTH - 2)
        h = random.randint(1, HEIGHT - 2)
        food = WIDTH * h + w
        cell_free = is_cell_free(food, snake_size, snake_array)

def make_move(pbest_move):
    """真正的蛇在这个函数中，朝pbest_move走1步"""
    global key, snake, board, snake_size, score, step_counter
    shift_array(snake, snake_size)
    snake[HEAD] += pbest_move
    p = snake[HEAD]

    draw_snake_and_food(snake, snake_size, food, score, high_score)
    step_counter += 1

    if snake[HEAD] == food:
        board[snake[HEAD]] = SNAKE
        snake_size += 1
        score += 1
        if snake_size < FIELD_SIZE:
            new_food()
    else:
        board[snake[HEAD]] = SNAKE
        board[snake[snake_size]] = UNDEFINED

def virtual_shortest_move():
    """虚拟运行吃到食物后，得到虚拟下蛇在board的位置"""
    global snake, board, snake_size, tmpsnake, tmpboard, tmpsnake_size, food
    tmpsnake_size = snake_size
    tmpsnake[:] = snake[:]
    tmpboard[:] = board[:]
    board_reset(tmpsnake, tmpsnake_size, tmpboard, tmpsnake[:tmpsnake_size], food)

    food_eated = False
    while not food_eated:
        board_BFS(food, tmpsnake, tmpboard, MOVES)
        move = choose_shortest_safe_move(tmpsnake, tmpboard, MOVES)
        shift_array(tmpsnake, tmpsnake_size)
        tmpsnake[HEAD] += move
        if tmpsnake[HEAD] == food:
            tmpsnake_size += 1
            board_reset(tmpsnake, tmpsnake_size, tmpboard, tmpsnake[:tmpsnake_size], food)
            tmpboard[food] = SNAKE
            food_eated = True
        else:
            tmpboard[tmpsnake[HEAD]] = SNAKE
            tmpboard[tmpsnake[tmpsnake_size]] = UNDEFINED

def find_safe_way():
    """如果蛇与食物间有路径，则调用本函数"""
    global snake, board
    safe_move = ERR
    virtual_shortest_move()
    if is_tail_inside():
        return choose_shortest_safe_move(snake, board, MOVES)
    safe_move = follow_tail()
    return safe_move

def game_loop():
    """游戏主循环"""
    global high_score, score, step_counter
    high_score = load_high_score()
    while True:
        board_reset(snake, snake_size, board, snake[:snake_size], food)
        if board_BFS(food, snake, board, MOVES):
            best_move = find_safe_way()
        else:
            best_move = follow_tail()
        if best_move == ERR:
            best_move = any_possible_move()
        if best_move != ERR:
            make_move(best_move)
        else:
            break
        if step_counter > MAX_STEPS:
            break
        draw_snake_and_food(snake, snake_size, food, score, high_score)
        fpsClock.tick(SNAKE_SPEED)  # 控制游戏速度

    # 检查是否需要更新最高分
    if score > high_score:
        high_score = score
        save_high_score(high_score)

    print("Game Over! Final Score:", score)